package com.fw.leetcode.heap;

import com.fw.leetcode.LeetCode;

/**
 * 215. Kth Largest Element in an Array
 *
 * Given an integer array nums and an integer k, return the kth largest element in the array.
 *
 * Note that it is the kth largest element in the sorted order, not the kth distinct element.
 *
 * Can you solve it without sorting? Only O(n) time complexity
 *
 * Example 1:
 *  Input: nums = [3,2,1,5,6,4], k = 2
 *  Output: 5
 *
 * Example 2:
 *  Input: nums = [3,2,3,1,2,4,5,5,6], k = 4
 *  Output: 4
 *
 * Constraints:
 *  1 <= k <= nums.length <= 10^5
 *  -10^4 <= nums[i] <= 10^4
 */
public class Num_0215 implements LeetCode {
    private interface Solution {
        int findKthLargest(int[] nums, int k);

        default void assertEquals(int[] nums, int k, int expected) {
            if (findKthLargest(nums, k) != expected) {
                findKthLargest(nums, k);
            }
        }
    }

    private static class MyQuickSolution implements Solution {

        @Override
        public int findKthLargest(int[] nums, int k) { // 快速排序变种（递归）：时O(n) 空O(log n)
            /*
             * 数组中第k个最大元素
             * - 不能用排序，还得 O(n) 时间复杂度
             * - 用快排分治思想，每次分治的`基准元素`位置可确定在最终排序的第几个位置，直到基准位置为k
             *                [5,1,3,8,2,4,9]
             * 第一次基准 5，left|right 双指针将<=5的放左边，大于5的放右边
             *  -(l|r-swap)-> [5,1,3,4,2,8,9] -(base-swap)-> [2,1,3,4,5,8,9]
             * 递归左右两部分并按上述方式按基准切分
             *
             * 用途：
             * - 解决 topN 问题
             */
            return quickSelect(nums, 0, nums.length - 1, nums.length - k);
        }

        int quickSelect(int[] nums, int start, int end, int k) {
            /*
             * 快速选择（升序）
             */
            if (start == end) return nums[k];
            int baseVal = nums[start]; // 左基准
            int left = start - 1; // 减1 保证 start==end 时仍能进入 while 循环
            int right = end + 1; // 加1 保证 start==end 时仍能进入 while 循环
            while (left < right) {
                do left++; while (nums[left] < baseVal); // 从左往右找比基准大的第一个
                do right--; while (nums[right] > baseVal); // 从右往左找比基准小的第一个
                if (left < right){
                    int tmp = nums[left];
                    nums[left] = nums[right];
                    nums[right] = tmp;
                }
            }
            if (k <= right) return quickSelect(nums, start, right, k);
            else return quickSelect(nums, right + 1, end, k);
        }
    }

    private static class MyHeapSolution implements Solution {

        @Override
        public int findKthLargest(int[] nums, int k) { // 其次：堆排序(非递归)：时O(nlog n) 空O(1)
            /*
             * 数组中第k个最大元素
             * - 不能用排序，还得 O(n) 时间复杂度
             * - 用堆（不满足 O(n) 时间复杂度）
             * 特性：
             * - 数组中元素构成了一颗满二叉树，不稳定排序
             *           0
             *     1           2
             *  3     4     5     6
             * 7 8   9 10 11 12 13 14
             * - 大顶堆：nums[i] >= max { nums[2*i+1], nums[2*i+2] }
             * - 小顶堆：nums[i] <= min { nums[2*i+1], nums[2*i+2] }
             * 用途：
             * - 解决 topN 问题
             */
            heapBuild(nums, 0, nums.length - 1); // 构建大根堆
            int result = -1;
            for (int i = 0; i < k; i++) {
                result = removeRoot(nums, 0, nums.length - 1 - i); // 移除堆顶（存于数组尾）
                heapify(nums, 0, nums.length - 2 - i); // 重新堆化（不包含最后移入的堆顶元素）
            }
            return result;
        }

        private void heapBuild(int[] nums, int start, int end) {
            /*
             * 大顶堆建立
             * - 从最后一棵子树（最后一个叶节点的父节点：(end-1)/2）开始向上遍历
             */
            for (int i = (end-1) / 2; i >= start; i--) {
                heapify(nums, i, end);
            }
        }

        private void heapify(int[] nums, int start, int end) { // 迭代（非递归）
            /*
             * 大顶堆堆化 nums[start, end]，使 nums[start] >= max { nums[start+1], ..., nums[end] }
             * 遍历路径：
             * - start -> 2*start+1 | 2*start+2 -> 到子树叶节点为止
             */
            int startNum = nums[start]; // 起始节点元素，（在最后交换位置一次性放入，避免多次交换）
            int i = 2 * start + 1; // 起始为左
            while (i <= end) {
                // 取出左右子树中较大的一个索引及元素，若只有左无右，则取左
                if (i + 1 <= end && nums[i] < nums[i+1]) { // 左与右比
                    i = i + 1;
                }
                // 子节点与起始节点比，若大则上移，否则结束
                if (startNum >= nums[i]) {
                    break;
                }
                nums[start] = nums[i]; // 子节点元素上移
                start = i; // 子树根索引下移
                i = 2 * start + 1; // 子树子节点索引下移
            }
            nums[start] = startNum; // 起始节点后移至最后子节点交换处
        }

        private int removeRoot(int[] nums, int start, int end) {
            /*
             * 移除大顶堆堆顶（最大值）
             * - 与最后一个元素交换位置
             */
            int top = nums[start];
            nums[start] = nums[end];
            nums[end] = top;
            return top;
        }
    }

    public static void main(String[] args) {
        Solution solution = new MyQuickSolution();
        solution.assertEquals(new int[]{1,2,3,4,5,6,7}, 5, 3);
        solution.assertEquals(new int[]{7,6,5,4,3,2,1}, 5, 3);
        solution.assertEquals(new int[]{3,3,3,3,3,3,3,3,3}, 1, 3);
        solution.assertEquals(new int[]{3,2,3,1,2,4,5,5,6}, 4, 4);
        solution.assertEquals(new int[]{2,1}, 2, 1);
        solution.assertEquals(new int[]{3,2,1,5,6,4}, 2, 5);
        int[] ints = new int[60000];
        for (int i = 0; i < 60000; i++) {
            ints[i] = 1;
        }
        solution.assertEquals(ints, 50000, 1);
    }
}
